package com.bins.langchain.easy.rag.service.impl;

import com.bins.langchain.easy.rag.main.Rag;
import com.bins.langchain.easy.rag.service.RagService;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.net.URISyntaxException;
import java.net.URL;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.List;

@Service
public class RagServiceImpl implements RagService {

    @Autowired
    private OllamaChatModel chatModel;

    @Autowired
    private EmbeddingStore<TextSegment> embeddingStore;

    @Override
    public void loadDocument() {
        URL url = Rag.class.getClassLoader().getResource("笑话.txt");
        Document document = getDocument(url);
        List<Document> documents = Arrays.asList(document);
        EmbeddingStoreIngestor.ingest(documents, embeddingStore);

        AiServices.builder(Assistant.class)
                .chatLanguageModel(chatModel)
                .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
                .build();
    }

    @Override
    public String chat(String prompt) {
        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(chatModel)
                .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
                .build();
        String answer = assistant.chat(prompt);
        System.out.println(answer);
        return answer;
    }

    public static Document getDocument(URL url) {
        Document document = null;
        try {
            Path path = Paths.get(url.toURI());
            document = FileSystemDocumentLoader.loadDocument(path);
        } catch (URISyntaxException e) {

        }
        return document;
    }
}

interface Assistant {

    String chat(String userMessage);
}
